import numpy as np
from matlab_flow import predict_flow

# 已知参数
p00 = 3.9040
p10 = 4.4143
p01 = -1.7234
p20 = -1.0054
p11 = 2.8593
p02 = -1.1549
p30 = 0.1765
p21 = -1.0957
p12 = 0.7047
p03 = 0.1436

# 归一化参数
mean_x = 1253
std_x = 383.6
mean_y = 234
std_y = 108.9


# 定义拟合函数
def predict_head(x, y):
    # 将输入值归一化
    x_norm = (x - mean_x) / std_x
    y_norm = (y - mean_y) / std_y

    # 多项式计算
    return (p00 +
            p10 * x_norm +
            p01 * y_norm +
            p20 * x_norm ** 2 +
            p11 * x_norm * y_norm +
            p02 * y_norm ** 2 +
            p30 * x_norm ** 3 +
            p21 * x_norm ** 2 * y_norm +
            p12 * x_norm * y_norm ** 2 +
            p03 * y_norm ** 3)


# 示例使用
x_input = 1140  # 输入的x值
y_input = 204.85  # 输入的y值

estimated_head = predict_head(x_input, y_input)
print(f'预估扬程: {estimated_head}')

# # 示例使用
# x_input = 1140  # 输入的x值
# y_input = 204.85  # 输入的y值
# estimated_flow = predict_flow(x_input, y_input)
# print(f'预估流量: {estimated_flow}')
